package com.example.dense_rack_server.controller.test;

import org.apache.pdfbox.cos.COSName;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.pdmodel.PDPage;
import org.apache.pdfbox.pdmodel.PDResources;
import org.apache.pdfbox.pdmodel.graphics.PDXObject;
import org.apache.pdfbox.pdmodel.graphics.image.PDImageXObject;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * pdf 图片提取器
 */

public class PdfImageExtractor {

    public static void main(String[] args) {
        // PDF文件路径
//        String pdfFilePath = "C:\\Users\\Administrator\\Desktop\\新建文件夹\\1.pdf";
//        String pdfFilePath = "C:\\Users\\Administrator\\Desktop\\新建文件夹\\13页.pdf";
        String pdfFilePath = "C:\\Users\\Administrator\\Desktop\\新建文件夹\\test.pdf";
        // 图片保存目录
        String outputDir = "C:\\Users\\Administrator\\Desktop\\新建文件夹\\";

        try {
            // 创建输出目录
            File outputDirectory = new File(outputDir);
            if (!outputDirectory.exists()) {
                outputDirectory.mkdirs();
            }

            // 加载PDF文档
            PDDocument document = PDDocument.load(new File(pdfFilePath));

            // 提取图片
            List<BufferedImage> images = extractImages(document);

            // 保存图片为PNG
            saveImagesAsPng(images, outputDir);

            // 关闭文档
            document.close();

            System.out.println("成功提取 " + images.size() + " 张图片");

        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 从PDF文档中提取所有图片
     */
    private static List<BufferedImage> extractImages(PDDocument document) throws IOException {
        List<BufferedImage> images = new ArrayList<>();

        // 遍历每一页
        for (PDPage page : document.getPages()) {
            PDResources resources = page.getResources();

            // 遍历所有资源
            for (COSName name : resources.getXObjectNames()) {
                PDXObject xObject = resources.getXObject(name);

                // 检查是否为图片
                if (xObject instanceof PDImageXObject) {
                    PDImageXObject image = (PDImageXObject) xObject;
                    BufferedImage bufferedImage = image.getImage();
                    images.add(bufferedImage);
                }
            }
        }

        return images;
    }

    /**
     * 将图片列表保存为PNG文件
     */
    private static void saveImagesAsPng(List<BufferedImage> images, String outputDir) throws IOException {
        for (int i = 0; i < images.size(); i++) {
            BufferedImage image = images.get(i);
            String fileName = outputDir + "image_" + (i + 1) + ".png";
            ImageIO.write(image, "png", new File(fileName));
            System.out.println("已保存: " + fileName);
        }
    }
}
